#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/execution_policy.h>
#include <thread>
// CUDA内核函数，使用thrust::stable_sort排序
struct is_not_zero
{
  __host__ __device__
  bool operator()(const int &x)
  {
    return x > 0;
  }
};

struct is_even
{
  __host__ __device__
  bool operator()(const int &x)
  {
    return (x % 2) == 0;
  }
};


int main() {
    int n = 10;  // 数据大小
    // thrust::host_vector<int> h_data(n);  // 主机端数据
    // thrust::device_vector<int> d_data;  // 设备端数据

    cudaDeviceProp device_props;
    cudaGetDeviceProperties(&device_props, 0);

    cudaSetDevice(0);


    int *h_vector_idx = new int[n];
    int *h_vector_score = new int[n];
    int *d_vector_idx = nullptr; cudaMalloc(&d_vector_idx, sizeof(int) * n);
    int *d_vector_score = nullptr; cudaMalloc(&d_vector_score, sizeof(int) * n);

    // 初始化数据
    for (int i = 0; i < n; i++) {
        h_vector_idx[i] = i % 2;
        h_vector_score[i] = i % 2;
    }
    
    // 将数据从主机端传输到设备端
    // d_data = h_data;

    cudaMemcpy(d_vector_idx, h_vector_idx, sizeof(int) * n, cudaMemcpyHostToDevice);
    cudaMemcpy(d_vector_score, h_vector_score, sizeof(int) * n, cudaMemcpyHostToDevice);
    clock_t time1,time2;


    int group = 2;
    time1 = clock();


    // thrust::partition(thrust::host, h_vector_score, h_vector_score + n, h_vector_idx, is_not_zero());
    thrust::partition(thrust::device, d_vector_score, d_vector_score + n, is_not_zero());
    
    // std::cout << end - d_vector_score << "\n";
    // std::cout << end - d_vector_score << "\n";


    time2 = clock();
    std::cout<<(double)(time2-time1)/CLOCKS_PER_SEC<<std::endl;
    
    // int A[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    // const int N = sizeof(A)/sizeof(int);
    // thrust::partition(thrust::host,
    //                 A, A + N,
    //                 is_even());
    
    for (int i = 0; i < n; i++) {
        std::cout << h_vector_idx[i] << " " << h_vector_score[i] << "\n";
    }

    // 将数据从设备端传输回主机端并打印结果
    cudaMemcpy(h_vector_idx, d_vector_idx, sizeof(int) * n, cudaMemcpyDeviceToHost);
    cudaMemcpy(h_vector_score, d_vector_score, sizeof(int) * n, cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();

    for (int i = 0; i < n; i++) {
        std::cout << h_vector_idx[i] << " " << h_vector_score[i] << "\n";
    }
    // thrust::copy(d_vector, d_vector + n, std::ostream_iterator<int>(std::cout, " "));
    // std::cout << std::endl;
    
    return 0;
}